// --------------------------------------------------------------------------------------------------------------------
// <copyright file="ConnectKinect.cs" company="Microsoft Corporation">
// The MIT License (MIT)
// 
// Copyright (c) 2014, Microsoft Corporation
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
//  all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
//  THE SOFTWARE.
// </copyright>
// --------------------------------------------------------------------------------------------------------------------
namespace Microsoft.Robotics.Vision.Cameras
{
    using System;
    using System.Collections.ObjectModel;
    using System.Runtime.InteropServices;
    using System.Threading.Tasks;

    using Microsoft.Kinect;
    using Microsoft.Robotics.Numerics;
    using Microsoft.Robotics.Tracing;

    /// <summary>
    /// Class to connect and initialize kinect
    /// </summary>
    public class ConnectKinect
    {
        /// <summary>
        /// Trace context for this class
        /// </summary>
        private static TraceContext connectKinectContext = TraceContexts.FromString(TraceContexts.Common, "ConnectKinect");

        /// <summary>
        /// The depthFrameReady delegate is assigned by the Kinect
        /// </summary>
        private Action<long, DepthImageStream, ImageFrameDepth, KinectCameraState> depthFrameReady;

        /// <summary>
        /// The colorFrameReady delegate is assigned by the Kinect
        /// </summary>
        private Action<long, ColorImageStream, ImageFrameRgb, KinectCameraState> colorFrameReady;

        /// <summary>
        /// The $irFrameReady$ delegate is assigned by the Kinect
        /// </summary>
        private Action<long, ColorImageStream, ImageFrame<ushort>, KinectCameraState> irFrameReady;

        /// <summary>
        /// The kinectSensor
        /// </summary>
        private KinectSensor kinectSensor;

        /// <summary>
        /// The depth image format to produce. 
        /// </summary>
        private DepthImageFormat? depthImageFormat;

        /// <summary>
        /// The RGB image format to produce. 
        /// </summary>
        private ColorImageFormat? rgbImageFormat;

        /// <summary>
        /// Kinect Depth Image Pixel struct which is used temporarily to copy the kinect image.
        /// </summary>
        private DepthImagePixel[] depthPixelData;

        /// <summary>
        /// Temporary buffer to copy raw Kinect RGB image data
        /// </summary>
        private byte[] kinectRgbPixelData;

        /// <summary>
        /// Disable correcting the kinect mirror imaging behavior for RGB images
        /// </summary>
        private bool leftRightImages;

        /// <summary>
        /// Force the infra red transmitter to be on by default
        /// Use a KinectCameraControlMessage to change it
        /// </summary>
        private bool forceInfraredEmitterOff = false;

        /// <summary>
        /// Initializes a new instance of the ConnectKinect class
        /// </summary>
        /// <param name="depthFrameReady">The depthImageReady Action delegate</param>
        /// <param name="colorFrameReady">The colorFrameReady Action delegate</param>
        /// <param name="irFrameReady">The $irFrameReady$ Action delegate</param>
        /// <param name="depthImageFormat">The desired frame format for the depth image. Undefined turns off depth image collection</param>
        /// <param name="rgbImageFormat">RGB image format. Undefined turns off RGB image collection</param>
        /// <param name="leftRightImages">Images produced are left right</param>
        public ConnectKinect(
            Action<long, DepthImageStream, ImageFrameDepth, KinectCameraState> depthFrameReady,
            Action<long, ColorImageStream, ImageFrameRgb, KinectCameraState> colorFrameReady,
            Action<long, ColorImageStream, ImageFrame<ushort>, KinectCameraState> irFrameReady,
            DepthImageFormat depthImageFormat,
            ColorImageFormat rgbImageFormat,
            bool leftRightImages)
        {
            KinectSensor.KinectSensors.StatusChanged += new EventHandler<StatusChangedEventArgs>(this.KinectStatusChange);

            this.depthFrameReady = depthFrameReady;
            this.colorFrameReady = colorFrameReady;
            this.irFrameReady = irFrameReady;
            this.depthImageFormat = depthImageFormat;
            this.rgbImageFormat = rgbImageFormat;
            this.leftRightImages = leftRightImages;
        }

        /// <summary>
        /// Gets RGB images byte count per pixel
        /// </summary>
        public static int RGBImageBytesPerPixel
        {
            get
            {
                return 3;
            }
        }

        /// <summary>
        /// Gets or sets the KinectSensor
        /// </summary>
        public KinectSensor KinectSensor
        {
            get
            {
                return this.kinectSensor;
            }

            set
            {
                this.kinectSensor = value;
            }
        }

        /// <summary>
        /// Detect and connect Kinect
        /// </summary>
        /// <returns>The alignment parameters for the kinect</returns>
        public ReadOnlyCollection<byte> DetectAndConnectKinect()
        {
            ReadOnlyCollection<byte> kinectParams = null;

            if (KinectSensor.KinectSensors.Count == 0)
            {
                throw new Exception("No Kinect sensors were found");
            }

            foreach (KinectSensor sensor in KinectSensor.KinectSensors)
            {
                if (sensor.Status == KinectStatus.Connected)
                {
                    this.KinectSensor = sensor;
                    break;
                }
                else
                {
                    Console.WriteLine("Kinect sensor status: {0}", sensor.Status.ToString());
                }
            }

            try
            {
                this.KinectSensor.ForceInfraredEmitterOff = this.forceInfraredEmitterOff;
            }
            catch (InvalidOperationException)
            {
                // Ok to ignore some hardware does not support this
                TraceOut.Warning(connectKinectContext, "Disabling infrared emitter is not supported by this hardware");
            }

            if (this.KinectSensor.Status == KinectStatus.Connected)
            {
                bool depthEnabled = false;
                bool rgbEnabled = false;

                // Undefined formats turns off image collection
                if (this.depthImageFormat != DepthImageFormat.Undefined)
                {
                    this.KinectSensor.DepthStream.Enable((DepthImageFormat)this.depthImageFormat);
                    depthEnabled = true;
                }

                if (this.rgbImageFormat != ColorImageFormat.Undefined)
                {
                    this.KinectSensor.ColorStream.Enable((ColorImageFormat)this.rgbImageFormat);
                    rgbEnabled = true;
                }

                // Start the kinect sensor
                this.KinectSensor.Start();

                // Get the alignment parameters
                kinectParams = new CoordinateMapper(this.KinectSensor).ColorToDepthRelationalParameters;

                if (true == depthEnabled)
                {
                    this.KinectSensor.DepthFrameReady += this.DepthImageReady;
                }

                if (true == rgbEnabled)
                {
                    this.KinectSensor.ColorFrameReady += this.RGBImageReady;
                }
            }

            return kinectParams;
        }

        /// <summary>
        /// DepthImageReady handler
        /// </summary>
        /// <param name="sender">The sender</param>
        /// <param name="e">The DepthImageFrameReadyEventArgs</param>
        public void DepthImageReady(object sender, DepthImageFrameReadyEventArgs e)
        {
            DepthImageStream depthStream = ((KinectSensor)sender).DepthStream;
            
            using (DepthImageFrame depthFrame = e.OpenDepthImageFrame())
            {
                if (depthFrame != null)
                {
                    ImageFrameDepth depthImage = this.PrepareDepthImage(depthFrame);

                    this.depthFrameReady(depthFrame.Timestamp, depthStream, depthImage, this.GetState());
                }
            }
        }

        /// <summary>
        /// RGB image handler
        /// </summary>
        /// <param name="sender">The sender</param>
        /// <param name="e">The EventArguments</param>
        public void RGBImageReady(object sender, ColorImageFrameReadyEventArgs e)
        {
            using (ColorImageFrame rgbFrame = e.OpenColorImageFrame())
            {
                if (rgbFrame != null)
                {
                    ColorImageStream rgbStream = ((KinectSensor)sender).ColorStream;

                    ImageFrameRgb rgbImage = null;
                    ImageFrame<ushort> irImage = null;

                    if (rgbFrame.Format == ColorImageFormat.InfraredResolution640x480Fps30)
                    {
                        irImage = this.PrepareIRImageFrame(rgbFrame);
                        this.irFrameReady(rgbFrame.Timestamp, rgbStream, irImage, this.GetState());
                    }
                    else
                    {
                        rgbImage = this.PrepareRGBImageFrame(rgbFrame);
                        this.colorFrameReady(rgbFrame.Timestamp, rgbStream, rgbImage, this.GetState());
                    }
                }
            }
        }

        /// <summary>
        /// KinectStatusChange handler
        /// </summary>
        /// <param name="sender">The sender</param>
        /// <param name="e">The DepthImageFrameReadyEventArgs</param>
        public void KinectStatusChange(object sender, StatusChangedEventArgs e)
        {
            if (this.KinectSensor == e.Sensor)
            {
                if (e.Status == KinectStatus.Disconnected || e.Status == KinectStatus.NotPowered)
                {
                    this.KinectSensor = null;
                    this.DetectAndConnectKinect();
                }
            }
        }

        /// <summary>
        /// Disconnect Kinect
        /// </summary>
        public void DisconnectKinect()
        {
            if (null != this.KinectSensor)
            {
                this.KinectSensor.Stop();
            }
        }

        /// <summary>
        /// Set the kinect camera state
        /// </summary>
        /// <param name="state">Requested kinect camera state</param>
        public void SetState(KinectCameraState state)
        {
            if (null != this.KinectSensor && KinectStatus.Connected == this.KinectSensor.Status)
            {
                try
                {
                    this.KinectSensor.ForceInfraredEmitterOff = state.EmitterOff;
                }
                catch (InvalidOperationException)
                {
                    // Ok to ignore some hardware does not support this
                    TraceOut.Warning(connectKinectContext, "Disabling infrared emitter is not supported by this hardware");
                }

                try
                {
                    if (state.NearModeOn == true)
                    {
                        this.KinectSensor.DepthStream.Range = DepthRange.Near;
                    }
                    else
                    {
                        this.KinectSensor.DepthStream.Range = DepthRange.Default;
                    }
                }
                catch (InvalidOperationException)
                {
                    // Ok to ignore some hardware does not support this
                    TraceOut.Warning(connectKinectContext, "Near mode not supported by this hardware");
                }
            }
        }

        /// <summary>
        /// Gets the state of the kinect camera
        /// </summary>
        /// <returns>State of camera</returns>
        public KinectCameraState GetState()
        {
            bool nearOff = (this.KinectSensor.DepthStream.Range == DepthRange.Default) ? false : true;
            KinectCameraState state = new KinectCameraState(this.KinectSensor.ForceInfraredEmitterOff, nearOff);
            return state;
        }

        /// <summary>
        /// Prepares a depth image frame from a kinect depth image frame. During preparation
        /// - No reading points are converted to the special NoReadingValue
        /// - When enabled depth buffer is aligned with image frame (NOTE for this to happen the
        ///   depth and image frames must be the same size)
        /// </summary>
        /// <param name="kinectDepthFrame">Kinect depth frame</param>
        /// <returns>Depth image frame</returns>
        private ImageFrameDepth PrepareDepthImage(DepthImageFrame kinectDepthFrame)
        {
            if (null == kinectDepthFrame)
            {
                return null;
            }

            ImageFrameDepth depthImage = new ImageFrameDepth(kinectDepthFrame.Width, kinectDepthFrame.Height);

            // reallocate memory if size has changed
            if (this.depthPixelData == null || this.depthPixelData.Length != kinectDepthFrame.PixelDataLength)
            {
                this.depthPixelData = new DepthImagePixel[kinectDepthFrame.PixelDataLength];
                }

            kinectDepthFrame.CopyDepthImagePixelDataTo(this.depthPixelData);
            unchecked
            {
                unsafe
                {
                    fixed (short* pixelData = depthImage.ImageData)
                    {
                        for (int i = 0; i < kinectDepthFrame.PixelDataLength; ++i)
                        {
                            if (this.leftRightImages)
                            {
                                int row = i / kinectDepthFrame.Width;
                                int col = i % kinectDepthFrame.Width;

                                pixelData[row * kinectDepthFrame.Width + (kinectDepthFrame.Width - 1 - col)] = this.depthPixelData[i].Depth;
                            }
                            else
                            {
                                pixelData[i] = this.depthPixelData[i].Depth;
                            }
                        }
                    }
                }
            }

            return depthImage;
        }

       /// <summary>
        /// Prepares a 2D image buffer as a simple byte array by converting the Kinect stream data
        /// </summary>
        /// <param name="image2DFrame">Kinect RGB image frame</param>
        /// <param name="doByteFlip">Flip the bytes</param>
        /// <returns>Converted byte buffer</returns>
        private byte[] Prepare2DImageBuffer(ColorImageFrame image2DFrame, bool doByteFlip)
        {
            int bytesPerPixel = Math.Min(image2DFrame.BytesPerPixel, ConnectKinect.RGBImageBytesPerPixel);
            byte[] rgbPixelData = null;

            if (null == image2DFrame)
            {
                return null;
            }

            if (this.kinectRgbPixelData == null || this.kinectRgbPixelData.Length != image2DFrame.PixelDataLength)
            {
                this.kinectRgbPixelData = new byte[image2DFrame.PixelDataLength];
            }

            image2DFrame.CopyPixelDataTo(this.kinectRgbPixelData);

            // Image needs to be transformed when the formats do not match or
            // mirror correction is required
                int pixRawIdx = 0;
                int bytePerPixelDifference = image2DFrame.BytesPerPixel - bytesPerPixel;
                int stride = image2DFrame.Width * bytesPerPixel;
                rgbPixelData = new byte[image2DFrame.Width * image2DFrame.Height * bytesPerPixel];

                // Use fixed/unchecked for perf reasons
                unsafe
                {
                    unchecked
                    {
                        fixed (byte* pixData = rgbPixelData)
                        {
                        fixed (byte* pixKinect = this.kinectRgbPixelData)
                            {
                                for (int row = 0; row < image2DFrame.Height; ++row)
                                {
                                int rowStart = row * image2DFrame.Width * bytesPerPixel;
                                if (this.leftRightImages)
                                {
                                    rowStart += (image2DFrame.Width - 1) * bytesPerPixel;
                                }

                                    int pixIdx;

                                    for (int col = 0; col < stride; col += bytesPerPixel)
                                    {
                                    if (false == this.leftRightImages)
                                        {
                                        pixIdx = rowStart + col;
                                        }
                                        else
                                        {
                                        pixIdx = rowStart - col;
                                        }

                                        int pixIdxTmp = pixIdx + bytesPerPixel - 1;

                                        if (false == doByteFlip)
                                        {
                                            for (int i = bytesPerPixel - 1; i >= 0; --i)
                                            {
                                                pixData[pixIdxTmp - i] = pixKinect[pixRawIdx];
                                                ++pixRawIdx;
                                            }
                                        }
                                        else
                                        {
                                            for (int i = 0; i < bytesPerPixel; ++i)
                                            {
                                                pixData[pixIdxTmp - i] = pixKinect[pixRawIdx];
                                                ++pixRawIdx;
                                            }
                                        }

                                        pixRawIdx += bytePerPixelDifference;
                                    }
                                }
                            }
                        }
                    }
            }

            return rgbPixelData;
        }

        /// <summary>
        /// Prepares a RGB camera frame
        /// </summary>
        /// <param name="kinectRgbFrame">RGB kinect frame acquired</param>
        /// <returns>RGB image frame</returns>
        private ImageFrameRgb PrepareRGBImageFrame(ColorImageFrame kinectRgbFrame)
        {
            byte[] rgbData = this.Prepare2DImageBuffer(kinectRgbFrame, true);

            ImageFrameRgb rgbImage = new ImageFrameRgb(kinectRgbFrame.Width, kinectRgbFrame.Height);
            unsafe
            {
                fixed (RgbPix* framePixPtr = rgbImage.ImageData)
                {
                    IntPtr ptrFrame = new IntPtr(framePixPtr);
                    Marshal.Copy(rgbData, 0, ptrFrame, rgbData.Length);
                }
            }

            return rgbImage;
        }

        /// <summary>
        /// Prepares a IR camera frame
        /// </summary>
        /// <param name="kinectRgbFrame">IR kinect frame acquired</param>
        /// <returns>IR image frame</returns>
        private ImageFrame<ushort> PrepareIRImageFrame(ColorImageFrame kinectRgbFrame)
        {
            int rgbBytesPerPixel = kinectRgbFrame.BytesPerPixel;
            byte[] irPixelData = this.Prepare2DImageBuffer(kinectRgbFrame, false);

            ImageFrame<ushort> irImage = new ImageFrame<ushort>(kinectRgbFrame.Width, kinectRgbFrame.Height);
            unsafe
            {
                fixed (ushort* pframe = irImage.ImageData)
                {
                    IntPtr ptrFrame = new IntPtr(pframe);
                    Marshal.Copy(irPixelData, 0, ptrFrame, irPixelData.Length);
                }
            }

            return irImage;
        }
    }
}
